import numpy as np
from gym.spaces import Box
import scipy.special

class CognitiveUAVEnv(object):
    def __init__(self, U=2, N=5, area_size=1000, H_u=40, P_max=5.0):
        # 系统参数
        self.U = U  # 无人机数量
        self.N = N  # 设备数量
        self.area_size = area_size  # 区域大小
        self.H_u = H_u  # 无人机固定飞行高度
        self.P_max = P_max  # 最大发射功率
        self.max_energy = 100000.0
        self.tau = 1.0  # 时隙长度
        self.trajectory = {u: [] for u in range(U)}# 轨迹
        self.Energy_params = {"Ps": 79.85, "Pm": 88.63, "Pk": 11.46, "Ur": 120, "Vh": 4.03, "d0": 0.6, "pa": 1.225,
                              "z": 0.05, "G": 0.503}

        # 初始化参数
        self._init_communication_params()
        self._init_mobility_params()

        # 强化学习空间
        self.action_space = self._build_action_space()
        self.observation_space = self._build_observation_space()


        # 状态初始化
        self.reset()

        # 能耗模型参数
        self.P_s = 79.85  # 悬停功率
        self.P_m = 88.63  # 最大连续功率
        self.P_k = 11.46  # 转子尖端功率
        self.U_r = 120.0  # 旋翼尖端速度
        self.V_h = 4.03  # 悬停诱导速度
        self.d_0 = 0.6  # 机身阻力比
        self.rho_a = 1.225  # 空气密度
        self.z = 0.05  # 旋翼实体比
        self.G = 0.503  # 旋翼盘面积

        # 归一化系数
        self.rate_norm = 1e6
        self.interference_norm = 30
        self.energy_norm = 100.0

    def _init_communication_params(self):
        """初始化通信相关参数"""
        self.a = 9.61  # LoS概率参数
        self.b = 0.61  # LoS概率参数
        # 路径损耗系数
        self.eta_LoS = 10 ** (-41.4 / 10)
        self.eta_NLoS = 10 ** (-60.4 / 10)
        self.sigma2 = 1e-10  # 噪声功率
        self.B = 4e6  # 带宽
        self.PT_power = 1  # 主发射机固定功率

    def _init_mobility_params(self):
        """初始化移动性参数"""
        self.v_max = 20.0
        self.theta_range = (-np.pi, np.pi)

    def _build_action_space(self):
        """构建动作空间"""

        action_dim = 2 * self.U + self.N + self.U * self.N
        return Box(low=-1, high=1, shape=(action_dim,), dtype=np.float32)

    def _build_observation_space(self):
        """构建状态空间"""

        state_dim = 1 + self.U * 3
        return Box(low=-1, high=1, shape=(state_dim,), dtype=np.float32)

    def reset(self):
        """重置环境状态"""
        self.T = 80  #任务周期
        self.current_time = 0


        # 初始化无人机位置
        self.uav_pos = np.array([[0.0, 1000.0, 40.0], [1000.0, 600.0, 40.0]])

        # 初始化设备位置
        self.device_pos = np.array([
            [200, 400, 0], [220, 300, 0], [300, 350, 0],
            [600, 300, 0], [700, 200, 0]
        ])

        # 主网络位置
        self.PR_pos = np.array([990, 990, 0])
        self.PT_pos = np.array([800, 800, 0])

        # 初始化设备关联矩阵
        self.alpha = np.zeros((self.U, self.N), dtype=int)
        self.energy_levels = np.full(self.U, self.max_energy)
        self.trajectory = {u: [] for u in range(self.U)}

        for u in range(self.U):
            self.trajectory[u].append(self.uav_pos[u].copy())  # 记录初始位置

        return self._get_state()

    def _get_state(self):

        state = []
        # 时间步归一化
        state.append(self.current_time / self.T * 2 - 1)

        # 无人机位置归一化
        uav_pos_norm = (self.uav_pos[:, :3] - np.array([self.area_size / 2, self.area_size / 2, self.H_u])) / \
                       np.array([self.area_size / 2, self.area_size / 2, self.H_u])
        state.extend(uav_pos_norm.flatten())

        return np.array(state, dtype=np.float32)

    def step(self, action, episode, t):
        """执行动作"""
        # 解析动作各部分
        ptr = 0
        # 1. 解析速度和角度
        velocities = (action[ptr:ptr + self.U] + 1) / 2 * self.v_max
        ptr += self.U
        angles = action[ptr:ptr + self.U] * np.pi
        ptr += self.U
        # 2. 解析设备分配
        device_actions = action[ptr:ptr + self.N]
        ptr += self.N
        self.alpha = np.zeros((self.U, self.N), dtype=int)
        for n in range(self.N):
            if device_actions[n] < 0:  # 小于0分配无人机0
                self.alpha[0, n] = 1
            else:  # 大于等于0分配无人机1
                self.alpha[1, n] = 1

        # 3. 解析功率分配
        power_actions = action[ptr:ptr + self.U * self.N].reshape(self.U, self.N)
        power_allocation = self._normalize_power(power_actions)


        self.prev_pos = self.uav_pos.copy()

        # 更新无人机位置
        boundary_violations = self._update_positions(velocities, angles)

        for u in range(self.U):
            self.trajectory[u].append(self.uav_pos[u].copy())

        # 计算性能指标
        rates = self._calculate_rates(power_allocation)
        energy_cost = self._calculate_energy_cost()
        self.energy_levels = np.clip(self.energy_levels - energy_cost, 0, None)
        pr_interference = self._calculate_pr_interference(power_allocation)

       #计算奖励
        reward = (
                (np.sum(rates)*abs(pr_interference) / self.rate_norm)/(np.sum(energy_cost))

                - (boundary_violations * 10)
        )
        # 终止条件
        done = (
                np.any(self.energy_levels <= 0)
                or (self.current_time >= self.T - 1)
        )

        info = {
            "rates": rates,
            "interference": pr_interference,
            "energy_used": energy_cost,
            "boundary_violations": boundary_violations,
            "alpha": self.alpha,
            "power_allocation": power_allocation
        }

        self.current_time += 1
        print("Episode:", episode, "t:", t, "Allocation:", self.alpha[0], self.alpha[1], "Location:", np.trunc(self.uav_pos[0]), np.trunc(self.uav_pos[1]), "reward:", reward, boundary_violations * 150, "|", np.sum(rates)/ self.rate_norm, (np.sum(rates) / self.rate_norm) * 45, rates/self.rate_norm)
        return self._get_state(), reward, done, info

    def _update_positions(self, velocities, angles):
        """更新无人机位置并返回边界违规次数"""
        violations = 0

        for u in range(self.U):
            dx = velocities[u] * np.cos(angles[u]) * self.tau
            dy = velocities[u] * np.sin(angles[u]) * self.tau

            self.uav_pos[u, 0] = self.uav_pos[u, 0] + dx
            self.uav_pos[u, 1] = self.uav_pos[u, 1] + dy

            if (self.uav_pos[u, 0] < 0) or (self.uav_pos[u, 0] > 1000) or (self.uav_pos[u, 1] < 0) or (self.uav_pos[u, 1] > 1000):
                violations += 1

        return violations

    def _calculate_path_loss(self, q_uav, q_device):
        """计算路径损耗"""
        d = np.linalg.norm(q_uav - q_device)
        if d == 0:
            return 0.0

        # 计算仰角（弧度）
        theta = np.arcsin(q_uav[2] / d) if d > 1e-6 else np.pi / 2
        theta_deg = np.degrees(theta)
        P_LoS = 1 / (1 + self.a * np.exp(-self.b * (theta_deg - self.a)))
        return (P_LoS * self.eta_LoS + (1 - P_LoS) * self.eta_NLoS) / (d ** 2)

    def _normalize_power(self, raw_power):
        """功率分配"""
        power_matrix = np.zeros((self.U, self.N))

        for u in range(self.U):
            # 获取当前无人机服务的设备
            served = [n for n in range(self.N) if self.alpha[u, n] == 1]
            if not served:
                continue

            # 提取该无人机的功率分配动作
            u_power = raw_power[u, served]

            # 按信道增益降序排序设备
            gains = []
            for n in served:
                # 计算无人机u到设备n的路径损耗
                h = self._calculate_path_loss(self.uav_pos[u], self.device_pos[n])
                # 计算外部干扰和噪声功率之和
                I_out = self._calculate_interference_power(u, n)
                # 归一化信道增益
                g = h / (I_out + self.sigma2)
                gains.append((n, g))
            sorted_devices = sorted(gains, key=lambda x: x[1], reverse=True)
            sorted_n = [n for (n, g) in sorted_devices]

            # 对功率动作进行Softmax
            power_weights = scipy.special.softmax(u_power)
            # 将 softmax 权重从小到大排序（大权重给弱用户）
            sorted_weights = np.sort(power_weights)

            # 按增益倒序分配递增功率

            for i, n in enumerate(sorted_n):
                idx_in_served = served.index(n)

                power_matrix[u, n] = (
                        sorted_weights[idx_in_served]

                        * self.P_max
                )

            # 总功率约束
            total_power = np.sum(power_matrix[u])
            if total_power > self.P_max:
                power_matrix[u] *= self.P_max / total_power
            #print(f"UAV{u} Power分配:", power_matrix[u][served])

        return power_matrix

    def _calculate_energy_cost(self):
        energy = np.zeros(self.U)
        if self.prev_pos is None: return energy

        for u in range(self.U):
            dx = self.uav_pos[u, 0] - self.prev_pos[u, 0]
            dy = self.uav_pos[u, 1] - self.prev_pos[u, 1]
            vh = np.hypot(dx, dy) / self.tau


            vv = 0

            energy[u] = self.Energy_params.get("Ps") * (
                        1 + 3 * (vh / self.Energy_params.get("Ur")) ** 2) + self.Energy_params.get(
                "Pm") * (np.sqrt(1 + (1 / 4) * (vh / self.Energy_params.get("Vh")) ** 4) - (1 / 2) * (
                    vh / self.Energy_params.get("Vh")) ** 2) ** (1 / 2) + (1 / 2) * self.Energy_params.get(
                "d0") * self.Energy_params.get("pa") * self.Energy_params.get("z") * self.Energy_params.get("G") * (
                        vh ** 3 + self.Energy_params.get("Pk") * vv)



        return energy

    def _calculate_rates(self, P_alloc):

        rates = np.zeros(self.N)

        for u in range(self.U):
            # 获取当前无人机 u 所服务的设备编号列表
            served = [n for n in range(self.N) if self.alpha[u, n] == 1]

            if not served:
                continue

            # 计算每个服务设备的增益并存储在 gains 列表中
            gains = []
            for n in served:
                # 计算无人机 u 到设备 n 的路径损耗
                h = self._calculate_path_loss(self.uav_pos[u], self.device_pos[n])
                # 计算设备 n 受到的干扰功率
                I = self._calculate_interference_power(u, n)
                # 计算增益并添加到 gains 列表中
                gains.append((n, h / (I + self.sigma2)))
            # 对 gains 列表按照增益值从大到小进行排序
            sorted_devices = sorted(gains, key=lambda x: x[1], reverse=True)

            # 按照 SIC 解码顺序计算每个设备的 SINR 和数据传输速率
            for i, (n, g) in enumerate(sorted_devices):
                # 计算区域内干扰

                intra = sum(P_alloc[u, sorted_devices[j][0]] *
                            self._calculate_path_loss(self.uav_pos[u], self.device_pos[n])
                            for j in range(i + 1, len(sorted_devices)))
                # 计算区域外干扰
                inter = self._calculate_interference_power(u, n)

                # 计算当前无人机 u 到设备 n 的路径损耗
                h_signal = self._calculate_path_loss(self.uav_pos[u], self.device_pos[n])
                # 计算信号功率
                signal = P_alloc[u, n] * h_signal
                # 噪声功率
                noise = self.sigma2

                # 根据 SIC 解码顺序计算 SINR
                if i == 0:
                    # 如果是第一个设备，则只考虑区域外干扰和噪声
                    SINR = signal / (inter + noise)

                else:
                    # 对于其他设备，需要考虑区域内干扰、区域外干扰和噪声
                    SINR = signal / (intra + inter + noise)

                # 根据香农公式计算设备 n 的数据传输速率
                rates[n] = self.B*np.log2(1 + SINR)


        return rates

    def _calculate_interference_power(self, u, n):
        """计算设备n受到的外部干扰"""
        interference = 0.0
        # 来自其他无人机的干扰
        for u_prime in range(self.U):
            if u_prime != u:
                h = self._calculate_path_loss(self.uav_pos[u_prime], self.device_pos[n])
                interference += h * self.P_max
        # 来自主发射机的干扰
        h_pt = self._calculate_path_loss(self.PT_pos, self.device_pos[n])
        interference += h_pt * self.PT_power
        return interference

    def _calculate_pr_interference(self, power_allocation=None):
        """计算对主接收机的总干扰"""
        if power_allocation is None:
            power_allocation = np.zeros((self.U, self.N))
        total = 0.0
        for u in range(self.U):
            h = self._calculate_path_loss(self.uav_pos[u], self.PR_pos)
            total += h * power_allocation[u].sum()


        total = max(total, 1e-12)
        return 10 * np.log10(total)

def main():
    # 初始化环境
    env = CognitiveUAVEnv()
    print("=== 环境测试开始 ===")
    print("观测空间维度:", env.observation_space.shape)
    print("动作空间维度:", env.action_space.shape)

    # 运行3个episode
    for episode in range(3):
        state = env.reset()
        done = False

        print(f"\n===== Episode {episode + 1} =====")
        for t in range(env.T):
            # 生成随机动作
            action = env.action_space.sample()

            # 执行动作
            next_state, reward, done, info = env.step(action, episode, t)

            # 解析关键信息
            total_rate = np.sum(info['rates']) / 1e6  # 转换为Mbps
            interference = info['interference']
            energy_used = np.sum(info['energy_used'])

            # 打印关键状态
            print(
                f"时隙 {t + 1:2d} |"
                f" 总速率: {total_rate:5.2f}Mbps |"
                f" 主网干扰: {interference:6.2f}dBW |"
                f" 能耗: {energy_used:6.2f}J |"
                f" 能量剩余: {env.energy_levels}"
            )

            # 检查终止条件
            if done:
                print(f"Episode 终止原因: {'能量耗尽' if any(env.energy_levels <= 0) else '时隙用尽'}")
                break


if __name__ == "__main__":
    main()
